import os
import glob
import xml.etree.ElementTree as ET
import argparse

def indent(elem, level=0):
    """ Helper function to indent the XML for pretty printing. """
    i = "\n" + level * "    "
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "    "
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            indent(elem, level + 1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i
        if not elem.tail:
            elem.tail = "\n"

def find_or_create_a_group(groups, group_name_target):
    # 找到指定 GroupName 的 Group 节点
    target_group = None
    for group in groups.findall('.//Group'):
        group_name = group.find('GroupName')
        if group_name is not None and group_name.text == group_name_target:
            target_group = group
            print(f"[GROUP] '{group_name_target}'")
            return target_group

    if target_group is None:
        target_group = ET.SubElement(groups, 'Group')
        target_group_name_node = ET.SubElement(target_group, 'GroupName')
        target_group_name_node.text = group_name_target
        print(f"[GROUP][+] '{group_name_target}'")
        return target_group


def add_include_path(uvprojx_file_path, new_include_path):
    tree = ET.parse(uvprojx_file_path)
    root = tree.getroot()
    Cads = root.find(".//Cads")
    include_path = Cads.find(".//IncludePath")

    if include_path is not None:
         # 获取 IncludePath 内容
        paths = include_path.text.split(';')
        paths = [path for path in paths if path]
        paths_rel = [os.path.relpath(p, start=os.path.dirname(uvprojx_file_path)) for p in paths]
        # 如果新路径不在已有路径中，则添加
        if new_include_path not in paths_rel:
            paths.append(new_include_path)
            include_path.text = ';'.join(paths) + ';'
            
            # 写回 XML 文件
            tree.write(uvprojx_file_path, encoding='utf-8', xml_declaration=True)
            print(f"Include Path '{new_include_path}' adds successfully.")

    else:
        print(f"'{uvprojx_file_path}' Maybe No a uvprojx")
        return


def update_files_to_group(uvprojx_file_path, folder_path):
    # 改变文件扩展名从 .uvprojx 到 .xml
    base, ext = os.path.splitext(uvprojx_file_path)
    if ext != '.uvprojx':
        print("The input file does not have the .uvprojx extension.")
        return

    if not os.path.exists(folder_path):
        print(f"File path '{folder_path}' does not exist.")
        return

    xml_path = base + '.xml'
    os.rename(uvprojx_file_path, xml_path)
    group_name_pre = os.path.basename(folder_path)

    try:
        # 解析XML文件
        tree = ET.parse(xml_path)
        root = tree.getroot()

        # 找到 Groups 节点
        groups = root.find('.//Groups');
        if groups is None:
            print(f"'{uvprojx_file_path}' Maybe No a uvprojx")
            return

        group_filepath_dict = {}
        group_filename_dict = {}
        # 找到目标 Groups 节点下的 Group 节点，如果不存在则创建一个
        for entry in os.listdir(folder_path):
            file_path = os.path.join(folder_path, entry)
            if os.path.isfile(file_path):
                # 查找所有 .c 和 .cpp 文件
                file = entry
                if 'base' not in group_filepath_dict:
                    group_filepath_dict['base'] = []
                    group_filename_dict['base'] = []
                if file.endswith('.c') or file.endswith('.cpp'):
                    group_filepath_dict['base'].append(os.path.relpath(file_path, start=os.path.dirname(xml_path)))
                    group_filename_dict['base'].append(os.path.basename(file_path))
            elif os.path.isdir(file_path):
                group_filepath_dict[entry] = []
                group_filename_dict[entry] = []
                # 遍历指定文件夹，查找所有 .c 和 .cpp 文件
                for subdir, _, files in os.walk(file_path):
                    for file in files:
                        if file.endswith('.c') or file.endswith('.cpp'):
                            sub_file_path = os.path.relpath(os.path.join(subdir, file), start=os.path.dirname(xml_path))
                            group_filepath_dict[entry].append(sub_file_path)
                            group_filename_dict[entry].append(os.path.basename(sub_file_path))

        # 删除文件夹中不存在的组
        for group in groups.findall('.//Group'):
            group_name = group.find('GroupName')
            group_name_str = group_name.text.split('/')
            if group_name is not None and len(group_name_str) >= 2 and group_name_str[0] == group_name_pre and group_name_str[1] not in group_filepath_dict:
                groups.remove(group)
                print(f"[GROUP][-] '{group_name.text}' is removed.")

        # 新增文件夹中存在的组
        for group in group_filepath_dict:
            group_name = os.path.basename(folder_path) + '/' + group
            target_group = find_or_create_a_group(groups, group_name)

            # 找到目标 Group 节点下的 Files 节点，如果不存在则创建一个
            files_node = target_group.find('Files')
            if files_node is None:
                files_node = ET.SubElement(target_group, 'Files')

            file_name_dict = {};
            for file in files_node.findall('.//File'):
                file_name = file.find('FileName').text
                file_path = file.find('FilePath').text
                if file_name not in group_filename_dict[group]:
                    files_node.remove(file)
                    print(f"    [FILE][-] '{file_name}' is removed.")
                else:
                    file_name_dict[file_name] = 1

            for i in range(len(group_filename_dict[group])):
                file_name = group_filename_dict[group][i]
                file_path = group_filepath_dict[group][i]

                if file_name in file_name_dict:
                    continue
                file_name_dict[file_name] = 1

                # 创建 File 节点并添加到 Files 节点下
                file_node = ET.SubElement(files_node, 'File')
                file_name_node = ET.SubElement(file_node, 'FileName')
                file_name_node.text = file_name
                file_type_node = ET.SubElement(file_node, 'FileType')
                if file_name.endswith('.c'):
                    file_type_node.text = '1'  # .c 文件类型都为 1
                elif file_name.endswith('.cpp'):
                    file_type_node.text = '8'  # .cpp 文件类型都为 8
                file_path_node = ET.SubElement(file_node, 'FilePath')
                file_path_node.text = file_path
                print(f"    [FILE][+] '{file_name}' is added.")

        # 格式化 XML
        indent(root)

        # 保存修改后的 XML 文件
        tree.write(xml_path, encoding='utf-8', xml_declaration=True)
        print(f"Update File Completed.")

    except ET.ParseError as e:
        print(f"ParseError: {e}")
        with open(xml_path, 'r', encoding='utf-8') as file:
            lines = file.readlines()
            start = max(0, e.position[0] - 5)
            end = min(len(lines), e.position[0] + 5)
            print("Context around the error:")
            for i in range(start, end):
                print(f"{i+1}: {lines[i].strip()}")

    finally:
        # 将文件扩展名改回 .uvprojx
        os.rename(xml_path, uvprojx_file_path)


def find_uvprojx_file():
    uvprojx_files = glob.glob("*.uvprojx")
    if not uvprojx_files:
        print("No .uvprojx file found in the current directory.")
        return None
    elif len(uvprojx_files) > 1:
        print("Multiple .uvprojx files found in the current directory:")
        for i, file in enumerate(uvprojx_files, start=1):
            print(f"{i}. {file}")
        print("Please ensure there is only one .uvprojx file in the directory.")
        return None
    else:
        return uvprojx_files[0]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Create the group and add .c and .cpp files in a .uvprojx file.")
    # parser.add_argument("folder_path", help="Relative path to the folder containing .c and .cpp files to be added.")
    # parser.add_argument("group_name", help="Name of the group to add the files to.")
    # args = parser.parse_args()

    uvprojx_file_path = find_uvprojx_file()
    if uvprojx_file_path:
        add_include_path(uvprojx_file_path, "..\\BSP")
        update_files_to_group(uvprojx_file_path, "..\\BSP")

    print(f"Mission Completed.")
    # os.system('pause')
